import fast_network_analysis as fna

import numpy

from scipy import *
import scipy.io as io

from brian import *

import shelve
import pickle
import time
import shutil
from datetime import datetime
import sys

class IFnetwork:

    #---------------- CONSTRUCTOR
    def __init__(self, Net_par,Syn_exc_par,Syn_inh_par,Syn_x_par,Pas_par,Fr_par,Init_par,Sim_par):
         self.Net_par = Net_par
         self.Syn_exc_par = Syn_exc_par
         self.Syn_inh_par = Syn_inh_par
         self.Syn_x_par = Syn_x_par
         self.Pas_par = Pas_par
         self.Fr_par = Fr_par
         self.Init_par = Init_par
         self.Sim_par = Sim_par
         



    #---------------- get all the parameters from the main file
    #---------------- prepare the neuron groups
    def buildPopulation(self):

        El=self.Pas_par['El'] * mvolt
        Ei=self.Pas_par['Ei'] * mvolt
        Ee=self.Pas_par['Ee'] * mvolt
        
        taum=self.Pas_par['C'] / self.Pas_par['gl'] *1000 *msecond

        tau_e_r=self.Syn_exc_par['tau'] * 1000 * msecond
        tau_e_d=self.Syn_exc_par['tau'] * 1000 * msecond

        tau_i_r=self.Syn_inh_par['tau'] * 1000 * msecond
        tau_i_d=self.Syn_inh_par['tau'] * 1000 * msecond

        eqs = Equations('''
dv/dt = ((El-v)+ge*(Ee-v)+gi*(Ei-v))*(1./taum) : mvolt
dge/dt = (x-ge)*(1./tau_e_d) : 1
dx/dt = -x*(1./tau_e_r) :1
dgi/dt = (y-gi)*(1./tau_i_d) : 1
dy/dt = -y*(1./tau_i_r) :1

''')


        self.Px=PoissonGroup(self.Net_par['Nx'],rates=self.Fr_par['x'] * Hz) # External drive
        
        self.P=NeuronGroup(self.Net_par['Ne']+self.Net_par['Ni'], # Recurrent network
                           model=eqs, 
                           threshold=self.Pas_par['thr'] * mvolt,
                           reset=self.Pas_par['thr_res'] * mvolt, 
                           refractory=self.Pas_par['tau_ref'] * msecond,
                           order=1, 
                           compile=True)

        self.Pe=self.P.subgroup(self.Net_par['Ne']) # Create excitatory and inhibitory popul
        self.Pi=self.P.subgroup(self.Net_par['Ni'])

    


    
    #---------------- make the connections
    def buildConnection(self):

        start_time = time.time()

        self.Cee = Connection(self.Pe, self.Pe, 'x') # create connections
        self.Cie = Connection(self.Pe, self.Pi, 'x')
        self.Cei = Connection(self.Pi, self.Pe, 'y')
        self.Cii = Connection(self.Pi, self.Pi, 'y')
        self.Cex = Connection(self.Px, self.Pe, 'x')
        self.Cix = Connection(self.Px, self.Pi, 'x')


        if self.Syn_exc_par['WeightDist'] == 'Normal': # fill them with weights
            self.Cee.connect_random(self.Pe,self.Pe,self.Net_par['Pee'],weight=lambda:numpy.random.normal(self.Syn_exc_par['Bee'],self.Syn_exc_par['Beestd']))

        elif self.Syn_exc_par['WeightDist'] == 'LogNormal':
              (lnmean,lnstd)=meanstd2lognorm(self.Syn_exc_par['Bee'],self.Syn_exc_par['Beestd'])

              self.Cee.connect_random(self.Pe,self.Pe,self.Net_par['Pee'],weight=lambda:numpy.random.lognormal(lnmean,lnstd))

        else: 
            print 'Unknown weight distribution\n'
            sys.exit(0)



        self.Cie.connect_random(self.Pe,self.Pi,
                                self.Net_par['Pie'],
                                weight=lambda:numpy.random.normal(self.Syn_exc_par['Bie'],self.Syn_exc_par['Biestd']))

        self.Cei.connect_random(self.Pi,self.Pe,
                                self.Net_par['Pei'],
                                weight=lambda:numpy.random.normal(self.Syn_inh_par['Bei'],self.Syn_inh_par['Beistd']))

        self.Cii.connect_random(self.Pi,self.Pi,
                                self.Net_par['Pii'],
                                weight=lambda:numpy.random.normal(self.Syn_inh_par['Bii'],self.Syn_inh_par['Biistd']))

        self.Cex.connect_random(self.Px,self.Pe,
                                self.Net_par['Pex'],
                                weight=lambda:numpy.random.normal(self.Syn_x_par['Bex'],self.Syn_x_par['Bexstd']))

        self.Cix.connect_random(self.Px,self.Pi,
                                self.Net_par['Pix'],
                                weight=lambda:numpy.random.normal(self.Syn_x_par['Bix'],self.Syn_x_par['Bixstd']))


        print "Network construction time:", time.time() - start_time, "seconds"


    

    #---------------- prepares the monitors
    def buildMonitors(self,n_rec_exc,n_rec_inh,lapse,time_step):
        
        temp=time.localtime(time.time())
        self.filename=str(temp[0])+'-'+str(temp[1])+'-'+str(temp[2])+'-'+str(temp[3])+'-'+str(temp[4])


        self.Me = PopulationSpikeCounter(self.Pe)
        self.Mi = PopulationSpikeCounter(self.Pi)
        self.Msp=SpikeMonitor(self.P)
        #self.Msp_aer=AERSpikeMonitor(self.P,'../datasave/'+self.filename+'.aedat')
        self.Mrates_exc = PopulationRateMonitor(self.Pe,0.01)
        self.Mrates_inh = PopulationRateMonitor(self.Pi,0.01)
        
        self.Mvolt=RecentStateMonitor(self.P,'v',duration=lapse,
                                record=arange(len(self.Pe)-n_rec_exc-1,len(self.Pe)+n_rec_inh),
                                timestep=time_step)
        self.Mge=RecentStateMonitor(self.P,'ge',duration=lapse,
                              record=arange(len(self.Pe)-n_rec_exc-1,len(self.Pe)+n_rec_inh),
                              timestep=time_step)
        self.Mgi=RecentStateMonitor(self.P,'gi',duration=lapse,
                              record=arange(len(self.Pe)-n_rec_exc-1,len(self.Pe)+n_rec_inh),
                              timestep=time_step)




    #---------------- calculate correlations (TO DELETE SOON)
    def correlationcalc(self):
        self.testcov=CCVF(self.Msp[3],self.Msp[3],width=40*ms,bin=1*ms,T=None)
        plot(self.testcov)
        show()


    #---------------- initialize the network
    def initialParam(self):
        vmean=self.Init_par['V'][0]
        vstd=self.Init_par['V'][1]

        gemean=self.Init_par['ge'][0]
        gestd=self.Init_par['ge'][1]
        
        gimean=self.Init_par['gi'][0]
        gistd=self.Init_par['gi'][1]


        self.P.v = (randn(len(self.P)) * vstd + vmean) * mvolt
        self.P.ge = randn(len(self.P)) * gestd + gemean
        self.P.gi = randn(len(self.P)) * gistd + gimean


    

    #---------------- run the simulations
    def runSim(self,duration,dosave='no'):

        

        self.net = Network(self.P,self.Px,
                      self.Cee,self.Cei,self.Cie,self.Cii,self.Cex,self.Cix,
                      self.Me,self.Mi,
                      self.Msp, 
                      #self.Msp_aer,
                      self.Mrates_exc,self.Mrates_inh,
                      self.Mvolt,self.Mge,self.Mgi)

        print "Simulation running..."
        start_time = time.time()

        self.net.run(duration*second,report='text',report_period=60.0*second)
        print "Simulation time:", time.time() - start_time, "seconds"

    


    #---------------- saves the spikemonitor in a dictionary and to file
    def saveSpikes(self,filename='def',folder='def'):
        if filename=='def':
            
            filename=self.filename
        
        if folder=='def':
            folder='../datasave/'

        pathfile=folder+filename

        s=shelve.open(pathfile+'_sp.db',writeback=True,protocol=-2)
        s['spikes']=self.Msp.spiketimes
        s.close()

        #output = open(pathfile+'.pkl', 'wb')
        #pickle.dump(self.Msp.spiketimes, output,-2)
        #output.close()
        
        #shutil.copyfile('main_class.py',pathfile+'_source.py')


    #---------------- saves the voltage to dictionaries and to file
    def saveMonitors(self,filename='def',folder='def'):

        if filename=='def':
          filename=self.filename

        if folder=='def':
            folder='../datasave/'

        pathfile=folder+filename
        
        s=shelve.open(pathfile+'_v'+'.db',writeback=True,protocol=-2)
        s['voltage']=self.Mvolt.values
        s.close()
        s=None

        f=shelve.open(pathfile+'_ge'+'.db',writeback=True,protocol=-2)
        f['ge']=self.Mge.values
        f.close()
        f=None

        g=shelve.open(pathfile+'_gi'+'.db',writeback=True,protocol=-2)
        g['gi']=self.Mgi.values
        g.close()

        g=None

        #shutil.copyfile('main_class.py',pathfile+'_source.py')
     

    #---------------- prints quickly the results to screen
    def grabResults(self):
        self.exc_data={}
        self.inh_data={}

        self.exc_data['Fr']=float(self.Me.nspikes) / self.Net_par['Ne']  / self.Sim_par['duration']
        self.inh_data['Fr']=float(self.Mi.nspikes) / self.Net_par['Ni']  / self.Sim_par['duration']    

        print "External Drive", self.Fr_par['x']
        print
        print "Firing rate"
        print "Excitatory mean", self.exc_data['Fr']
        print "Inhibitory mean", self.inh_data['Fr']
        print
        print "Membrane potential"
        Mean_mean_vm, Std_mean_vm, Mean_std_vm, Std_std_vm = fna.calc_stats_monitor(self.Mvolt,self.Sim_par['Monitors_exc_size'],self.Sim_par['Monitors_inh_size'])
        print
        print "Ge"
        Mean_mean_ge, Std_mean_ge, Mean_std_ge, Std_std_ge = fna.calc_stats_monitor(self.Mge,self.Sim_par['Monitors_exc_size'],self.Sim_par['Monitors_inh_size'])
        print
        print "Gi"
        Mean_mean_gi, Std_mean_gi, Mean_std_gi, Std_std_gi = fna.calc_stats_monitor(self.Mgi,self.Sim_par['Monitors_exc_size'],self.Sim_par['Monitors_inh_size'])

        self.exc_data['Vmmean']=Mean_mean_vm['exc']
        self.exc_data['Vmstd']=Mean_std_vm['exc']
        self.exc_data['Ge']=Mean_mean_ge['exc']
        self.exc_data['Gi']=Mean_mean_gi['exc']
        self.exc_data['Reff']=((Mean_mean_gi['exc']+Mean_mean_ge['exc']+1)*self.Pas_par['gl'])**(-1)
        self.exc_data['Taueff']=self.exc_data['Reff']*self.Pas_par['C']


        self.inh_data['Vmmean']=Mean_mean_vm['inh']
        self.inh_data['Vmstd']=Mean_std_vm['inh']
        self.inh_data['Ge']=Mean_mean_ge['inh']
        self.inh_data['Gi']=Mean_mean_gi['inh']
        self.inh_data['Reff']=((Mean_mean_gi['inh']+Mean_mean_ge['inh']+1)*self.Pas_par['gl'])**(-1)
        self.inh_data['Taueff']=self.inh_data['Reff']*self.Pas_par['C']


   
    #---------------- resets the network to start new simulation
    def resetNet(self):
        self.net.reinit(states=True)
        self.initialParam()



    #---------------- parameter loop
    def loop1D(self,parname,minval,maxval,nsteps):
        
        self.val_loop1D={}
        self.val_loop1D['exc']={}
        self.val_loop1D['exc']['Fr']=[]
        self.val_loop1D['exc']['Reff']=[]
        self.val_loop1D['exc']['Taueff']=[]
        self.val_loop1D['exc']['Ge']=[]
        self.val_loop1D['exc']['Gi']=[]
        self.val_loop1D['exc']['Vmmean']=[]
        self.val_loop1D['exc']['Vmstd']=[]


        self.val_loop1D['inh']={}
        self.val_loop1D['inh']['Fr']=[]
        self.val_loop1D['inh']['Reff']=[]
        self.val_loop1D['inh']['Taueff']=[]
        self.val_loop1D['inh']['Ge']=[]
        self.val_loop1D['inh']['Gi']=[]
        self.val_loop1D['inh']['Vmmean']=[]
        self.val_loop1D['inh']['Vmstd']=[]


        par_list=linspace(minval,maxval,num=nsteps)
        self.val_loop1D['par_val']=par_list

        for x in par_list:
            
            self.resetNet()
            self.Fr_par['x']=x
            
            self.buildPopulation()
            self.buildConnection()
            self.buildMonitors(self.Sim_par['Monitors_exc_size'],self.Sim_par['Monitors_inh_size'])
            self.initialParam()
            self.runSim(self.Sim_par['duration'])
            self.grabResults()
            
            self.val_loop1D['exc']['Fr'].append(self.exc_data['Fr'])
            self.val_loop1D['exc']['Reff'].append(self.exc_data['Reff'])
            self.val_loop1D['exc']['Taueff'].append(self.exc_data['Taueff'])
            self.val_loop1D['exc']['Ge'].append(self.exc_data['Ge'])
            self.val_loop1D['exc']['Gi'].append(self.exc_data['Gi'])
            self.val_loop1D['exc']['Vmmean'].append(self.exc_data['Vmmean'])
            self.val_loop1D['exc']['Vmstd'].append(self.exc_data['Vmstd'])

            self.val_loop1D['inh']['Fr'].append(self.inh_data['Fr'])
            self.val_loop1D['inh']['Reff'].append(self.inh_data['Reff'])
            self.val_loop1D['inh']['Taueff'].append(self.inh_data['Taueff'])
            self.val_loop1D['inh']['Ge'].append(self.inh_data['Ge'])
            self.val_loop1D['inh']['Gi'].append(self.inh_data['Gi'])
            self.val_loop1D['inh']['Vmmean'].append(self.inh_data['Vmmean'])
            self.val_loop1D['inh']['Vmstd'].append(self.inh_data['Vmstd'])


#---------------- save parameters
    def save_loop(self,filename='def',folder='def'):
        if filename=='def':
            filename=str(datetime.now())
        
        if folder=='def':
            folder='../datasave/'

        pathfile=folder+filename
        import scipy.io as io
        
        io.savemat(pathfile+'_loop',self.val_loop1D)
        
        shutil.copyfile('main_class.py',pathfile+'_source.py')



# -------------- FUNCTION USED BY THE CLASS

def meanstd2lognorm(a,b):
    aln=log((a**2)/(sqrt((b)**2 + a**2)))
    bln=sqrt(log((b)**2/(a**2)+1))

    return aln, bln
